#include <stdlib.h>
#include <stdio.h>
#include <signal.h>
#include <arpa/inet.h>
#include <netlink/netlink.h>
#include <netlink/genl/genl.h>
#include <netlink/genl/ctrl.h>

#include "../modules/ivi_config.h"
#include "../modules/ivi_nl.h"

#ifndef NIP4
#define NIP4(addr) \
	((unsigned char *)&addr)[3], \
	((unsigned char *)&addr)[2], \
	((unsigned char *)&addr)[1], \
	((unsigned char *)&addr)[0]
#define NIP4_FMT "%u.%u.%u.%u"
#endif

#ifndef NIP6
#define NIP6(addr) \
	ntohs((addr).s6_addr16[0]), \
	ntohs((addr).s6_addr16[1]), \
	ntohs((addr).s6_addr16[2]), \
	ntohs((addr).s6_addr16[3]), \
	ntohs((addr).s6_addr16[4]), \
	ntohs((addr).s6_addr16[5]), \
	ntohs((addr).s6_addr16[6]), \
	ntohs((addr).s6_addr16[7])
#define NIP6_FMT "%04x:%04x:%04x:%04x:%04x:%04x:%04x:%04x"
#endif

#define NLMSG_SIZE 1024

static struct nl_sock *sk = NULL;
static int family = 0, group = 0;

static int sd;
static struct sockaddr_in ms_addr;

static char *dev_name = "eth0";

static struct nla_policy ivi_genl_policy[IVI_NLA_MAX + 1] = {
	[IVI_NLA_V4ADDR] = { .type = NLA_U32 },
	[IVI_NLA_PLEN4]	= { .type = NLA_U32 },
	[IVI_NLA_V6ADDR] = { .type = NLA_UNSPEC, .minlen = sizeof(struct in6_addr) },
	[IVI_NLA_PLEN6] = { .type = NLA_U32 },
	[IVI_NLA_RATIO] = { .type = NLA_U16 },
	[IVI_NLA_OFFSET] = { .type = NLA_U16 },
	[IVI_NLA_FORMAT] = { .type = NLA_U8 },
	[IVI_NLA_ROUTE] = { .type = NLA_UNSPEC, .minlen = sizeof(struct in6_addr) },
};


#define MSG_T_CLIENT 0
#define MSG_T_SERVER 1

#define MSG_C_LOOKUP 0
#define MSG_C_UPDATE 1

struct lookup_request_msg {
    unsigned char type;
    unsigned char code;
    unsigned short rsvd1;
    unsigned long seq;
    unsigned long v4addr;
    unsigned short rsvd2;
    unsigned short offset;
};

struct update_request_msg {
    unsigned char type;
    unsigned char code;
    unsigned short rsvd;
    unsigned long seq;
    unsigned long v4addr;
    unsigned short ratio;
    unsigned short offset;
    struct in6_addr v6prefix;
    unsigned short plen6;
    unsigned char key[2];
};

struct lookup_reply_msg {
    unsigned char type;
    unsigned char code;
    unsigned short rsvd1;
    unsigned long seq;
    unsigned long v4addr;
    unsigned short ratio;
    unsigned short offset;
    struct in6_addr v6prefix;
    unsigned short plen6;
    unsigned short rsvd2;
};

static unsigned char rsa_key[3] = { 7, 23, 187 };  // e, d, n

static inline int fls(int ratio) {
	int ret;
	switch (ratio) {
		case 4096: ret = 13; break;
		case 2048: ret = 12; break;
		case 1024: ret = 11; break;
		case 512: ret = 10; break;
		case 256: ret = 9; break;
		case 128: ret = 8; break;
		case 64: ret = 7; break;
		case 32: ret = 6; break;
		case 16: ret = 5; break;
		case 8: ret = 4; break;
		case 4: ret = 3; break;
		case 3: ret = 2; break;
		case 2: ret = 1; break;
		case 1:
		default:
			ret = 0; break;
	}
	return ret;
}

static inline unsigned char rsa_compute(unsigned char c, unsigned char d, unsigned char n) {
	int i, tmp = 1;
	// Compute c^d % n
	for (i = 0; i < d; i++) {
		tmp *= c;
		tmp %= n;
	}
	return tmp;
}

static int ivid_callback_dispatch(struct nl_msg *msg, void *arg) {
	struct nlmsghdr *nlh = nlmsg_hdr(msg);
	struct genlmsghdr *gnlh = (struct genlmsghdr *)nlmsg_data(nlh);
	struct nlattr *attrs[IVI_NLA_MAX + 1];

	genlmsg_parse(nlh, 0, attrs, IVI_NLA_MAX, ivi_genl_policy);

	switch (gnlh->cmd) {
		case IVI_NLC_CLOSE:
			printf("ivid: closed by kernel.\n");
			exit(0);
			break;

		case IVI_NLC_RULE:
			if (fork() == 0) {
				//printf("ivid_callback_dispatch: child executes...\n");
				if (attrs[IVI_NLA_V4ADDR] && attrs[IVI_NLA_OFFSET]) {
					unsigned long addr = nla_get_u32(attrs[IVI_NLA_V4ADDR]);
					unsigned short port = nla_get_u16(attrs[IVI_NLA_OFFSET]);
					printf("ivid_callback_dispatch: kernel asks for " NIP4_FMT ":%d on child pid %d\n", NIP4(addr), port, getpid());
					
					socklen_t ms_len = sizeof(ms_addr);
					int recv_len;
					int seq = rand();
					char buffer[64];
					struct lookup_request_msg *lookup = (struct lookup_request_msg *)buffer;
					lookup->type = MSG_T_CLIENT;
					lookup->code = MSG_C_LOOKUP;
					lookup->rsvd1 = 0;
					lookup->seq = seq;
					lookup->v4addr = htonl(addr);
					lookup->rsvd2 = 0;
					lookup->offset = htons(port);

					//printf("ivid_callback_dispatch: seq %d on child pid %d\n", seq, getpid());
					
					sendto(sd, buffer, sizeof(struct lookup_request_msg), 0, (struct sockaddr *)&ms_addr, ms_len);
					recv_len = recvfrom(sd, buffer, 100, 0, (struct sockaddr *)&ms_addr, &ms_len);
					if (recv_len < 0) {
						printf("ivid_callback_dispatch: recvfrom() failed on child pid %d\n", getpid());
						exit(0);
					}

					printf("ivid_callback_dispatch: receive server reply on child pid %d\n", getpid());
					//printf("ivid_callback_dispatch: child wake up from pid %d\n", getpid());

					struct lookup_reply_msg *rlookup = (struct lookup_reply_msg *)buffer;
					if (rlookup->type != MSG_T_SERVER || rlookup->code != MSG_C_LOOKUP) {
						printf("ivid_callback_dispatch: message type unmatch\n");
						exit(0);
					}
					if (rlookup->seq != seq) {
						printf("ivid_callback_dispatch: wrong seq %d received on child pid %d\n", rlookup->seq, getpid());
						//exit(0);
					}
					
					rlookup->v4addr = ntohl(rlookup->v4addr);
					rlookup->ratio = ntohs(rlookup->ratio);
					rlookup->offset = ntohs(rlookup->offset);
					rlookup->plen6 = ntohs(rlookup->plen6);
					printf("ivid_callback_dispatch: " NIP4_FMT " r=%d o=%d -> " NIP6_FMT "/%d\n", 
						NIP4(rlookup->v4addr), rlookup->ratio, rlookup->offset, NIP6(rlookup->v6prefix), rlookup->plen6);
					
					struct nl_msg *rmsg;
					rmsg = nlmsg_alloc_size(NLMSG_SIZE);
					genlmsg_put(rmsg, NL_AUTO_PID, NL_AUTO_SEQ, family, 0, 0, IVI_NLC_RULE, IVI_NL_VERSION);
					nla_put_u32(rmsg, IVI_NLA_V4ADDR, rlookup->v4addr);
					nla_put(rmsg, IVI_NLA_V6ADDR, sizeof(struct in6_addr), &rlookup->v6prefix);
					nla_put_u32(rmsg, IVI_NLA_PLEN6, rlookup->plen6);
					nla_put_u16(rmsg, IVI_NLA_RATIO, rlookup->ratio);
					nla_put_u16(rmsg, IVI_NLA_OFFSET, rlookup->offset);
					nla_put_u8(rmsg, IVI_NLA_FORMAT, ADDR_FMT_SUFFIX);

					nl_send_auto_complete(sk, rmsg);
					nlmsg_free(rmsg);
				}
				//printf("ivid_callback_dispatch: child exit...\n");
				exit(0);
			}
			//printf("ivid_callback_dispatch: parent return...\n");
			break;

		case IVI_NLC_UPDATE:
			if (fork() == 0) {
				if (attrs[IVI_NLA_V4ADDR] && attrs[IVI_NLA_RATIO] && attrs[IVI_NLA_OFFSET] 
					&& attrs[IVI_NLA_V6ADDR] && attrs[IVI_NLA_PLEN6]) {
					unsigned long addr = nla_get_u32(attrs[IVI_NLA_V4ADDR]);  // in host byte order
					unsigned short ratio = nla_get_u16(attrs[IVI_NLA_RATIO]);
					unsigned short offset = nla_get_u16(attrs[IVI_NLA_OFFSET]);
					struct in6_addr *v6addr = nla_data(attrs[IVI_NLA_V6ADDR]);
					int plen6 = nla_get_u32(attrs[IVI_NLA_PLEN6]);
			
					printf("ivid_callback_dispatch: kernel updates " NIP4_FMT " r=%d o=%d -> " NIP6_FMT "/%d\n", 
						NIP4(addr), ratio, offset, NIP6(*v6addr), plen6);

					if (attrs[IVI_NLA_ROUTE]) {
						// Update interface address and route first
						printf("ivid_callback_dispatch: kernel requests to reconfig v6 address and route.\n");
						struct in6_addr naddr;
						struct in6_addr *router = nla_data(attrs[IVI_NLA_ROUTE]);
						char cmd[128];
						int prefixlen = plen6 >> 3;

						// Config address
						memcpy(&naddr, v6addr, sizeof(naddr));
						naddr.s6_addr[prefixlen] = (unsigned char)(addr >> 24);
						naddr.s6_addr[prefixlen + 1] = (unsigned char)((addr >> 16) & 0xff);
						naddr.s6_addr[prefixlen + 2] = (unsigned char)((addr >> 8) & 0xff);
						naddr.s6_addr[prefixlen + 3] = (unsigned char)(addr & 0xff);

						if (ratio != 0 && ratio != 1) {
							unsigned short suffix = fls(ratio) - 1;
							suffix = suffix << 12;
							suffix += offset & 0x0fff;
							naddr.s6_addr[prefixlen + 4] = (suffix >> 8) & 0xff;
							naddr.s6_addr[prefixlen + 5] = suffix & 0xff;
						}

						// Ready...
						sprintf(cmd, "ip -6 addr flush dev %s scope global", dev_name);
						printf("ivid_callback_dispatch: build addr cmd1 '%s' \n", cmd);
						// Action!
						system(cmd);
						usleep(1000);  // Wait for the system call to complete

						// Ready...
						sprintf(cmd, "ip -6 addr add " NIP6_FMT "/%d dev %s", NIP6(naddr), plen6, dev_name);
						printf("ivid_callback_dispatch: build addr cmd2 '%s' \n", cmd);
						// Action!
						system(cmd);
						usleep(1000);  // Wait for the system call to complete

						// Config route
						system("ip -6 route flush default");
						usleep(1000);  // Wait for the system call to complete

						// Ready...
						sprintf(cmd, "ip -6 route add default via " NIP6_FMT " dev eth0", NIP6(*router));
						printf("ivid_callback_dispatch: build route cmd2 '%s' \n", cmd);
						// Action!
						system(cmd);
						usleep(1000);  // Wait for the system call to complete
					}
					
					socklen_t ms_len = sizeof(ms_addr);
					int recv_len;
					int seq = rand();
					char buffer[64];
					struct update_request_msg *update = (struct update_request_msg *)buffer;
					update->type = MSG_T_CLIENT;
					update->code = MSG_C_UPDATE;
					update->rsvd = 0;
					update->seq = seq;
					update->v4addr = htonl(addr);
					update->ratio = htons(ratio);
					update->offset = htons(offset);
					memcpy(&update->v6prefix, v6addr, sizeof(struct in6_addr));
					update->plen6 = htons((unsigned short)plen6);
					update->key[0] = rsa_key[2];
					update->key[1] = rsa_compute((seq & 0xff) % rsa_key[2], rsa_key[1], rsa_key[2]);  // Use last byte in seq mod n as plaintext

					//printf("ivid_callback_dispatch: seq %d on child pid %d\n", seq, getpid());
					
					if (sendto(sd, buffer, sizeof(struct update_request_msg), 0, (struct sockaddr *)&ms_addr, ms_len) < 0) {
						printf("ivid_callback_dispatch: sendto() failed when updating on child pid %d\n", getpid());
					}
					printf("ivid_callback_dispatch: sendto() success on child pid %d\n", getpid());
				}
				exit(0);
			}
			break;
			
		default:
			/* ignore unknown cmd */
			printf("ivid_callback_dispatch: i don't understand cmd %d\n", gnlh->cmd);
			break;
	}
	
	return 0;
}

int main() {
	sk = nl_socket_alloc();

	nl_socket_disable_seq_check(sk);

	nl_socket_modify_cb(sk, NL_CB_VALID, NL_CB_CUSTOM, ivid_callback_dispatch, (void*)sk);

	if (genl_connect(sk)) {
		printf("ivid: genl_connect() failed.\n");
		exit(-1);
	}

	family = genl_ctrl_resolve(sk, IVI_FAMILY_NAME);
	if (family < 0) {
		printf("ivid: genl_ctrl_resolve() failed.\n");
		exit(-1);
	}

	group = genl_ctrl_resolve_grp(sk, IVI_FAMILY_NAME, IVI_GROUP_NAME);
	if (group < 0) {
		printf("ivid: genl_ctrl_resolve_grp() failed.\n");
		exit(-1);
	}

	if (nl_socket_add_membership(sk, group) < 0) {
		printf("ivid: nl_socket_add_membership() failed.\n");
		exit(-1);
	}

	if ((sd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
		printf("ivid: socket() failed.\n");
		exit(-1);
	}

	struct timeval tv;
	tv.tv_sec = 10;
	tv.tv_usec = 0;
	setsockopt(sd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

	memset(&ms_addr, 0, sizeof(ms_addr));
	ms_addr.sin_family = AF_INET;
	if (inet_pton(AF_INET, "2.2.2.1", &ms_addr.sin_addr) <= 0) {
		printf("ivid: inet_pton() failed.\n");
		close(sd);
		exit(-1);
	}
	ms_addr.sin_port = htons(5001);
	
/*
	if (daemon(1,0) != 0) {
		printf("ivid: daemon() failed.\n");
		exit(-1);
	}
*/
	signal(SIGCHLD, SIG_IGN);

	printf("ivid: so begins the story...\n");

	while (1) {
		if (nl_recvmsgs_default(sk) < 0) {
			printf("ivid: nl_recvmsgs_default() failed.\n");
			break;
		}
	}

	return 0;
}
